open Ast

exception Semantic_error of string

module StringMap = Map.Make(String)

type var_info = {
  vtype : ty;
  defined_at : Lexing.position option;
}

type func_info = {
  return_type : ty;
  params : param list;
  defined_at : Lexing.position option;
  has_return : bool; (* 添加是否包含return语句的标记 *)
}

(* 修改环境结构以支持作用域栈 *)
type scope = {
  vars : var_info StringMap.t;
}

type environment = {
  scopes : scope list;  (* 作用域栈，当前作用域在头部 *)
  funcs : func_info StringMap.t;
  in_loop : bool;
  current_func : string option; (* 添加当前函数名记录 *)
}

let empty_env = {
  scopes = [{ vars = StringMap.empty }];
  funcs = StringMap.empty;
  in_loop = false;
  current_func = None;
}

(* 进入新作用域 *)
let enter_scope env = 
  { env with scopes = { vars = StringMap.empty } :: env.scopes }

(* 离开当前作用域 *)
let leave_scope env = 
  match env.scopes with
  | [] -> failwith "Scope stack underflow"
  | _ :: scopes -> { env with scopes = scopes }

(* 查找变量（遍历所有作用域） *)
let rec find_var env id =
  match env.scopes with
  | [] -> None
  | scope :: rest ->
      match StringMap.find_opt id scope.vars with
      | Some info -> Some info
      | None -> find_var { env with scopes = rest } id

(* 在当前作用域添加变量 *)
let add_var env id info =
  match env.scopes with
  | [] -> failwith "No current scope"
  | scope :: rest ->
      let new_scope = { vars = StringMap.add id info scope.vars } in
      { env with scopes = new_scope :: rest }

let rec check_expr env expr = 
  match expr with
  | Number _ -> Int
  | Identifier id -> 
      (match find_var env id with
      | None -> raise (Semantic_error ("Undefined variable: " ^ id))
      | Some info -> info.vtype)
  | UnaryOp (_, e) -> check_expr env e
  | BinaryOp (op, e1, e2) ->
      let t1 = check_expr env e1 in
      let t2 = check_expr env e2 in
      if t1 != Int || t2 != Int then
        raise (Semantic_error "Operands must be integers");
      (match op with
      | Add | Sub | Mul | Div | Mod | Lt | Gt | Le | Ge -> Int
      | Eq | Ne | And | Or -> Int
      | Assign -> t1)
  | Call (fname, args) ->
      if not (StringMap.mem fname env.funcs) then
        raise (Semantic_error ("Undefined function: " ^ fname));
      let func = StringMap.find fname env.funcs in
      if List.length args != List.length func.params then
        raise (Semantic_error "Incorrect number of arguments");
      List.iter2 (fun arg (ptype, _) ->
        let atype = check_expr env arg in
        if atype != ptype then
          raise (Semantic_error "Argument type mismatch")
      ) args func.params;
      func.return_type
  | Paren e -> check_expr env e

(* 修改：添加return语句标记 *)
let rec check_stmt env stmt return_type = 
  match stmt with
  | Block stmts -> 
      let env = enter_scope env in
      let env, has_return = 
        List.fold_left (fun (env, has_ret) stmt -> 
          let env, ret = check_stmt env stmt return_type in
          (env, has_ret || ret)
        ) (env, false) stmts 
      in
      (leave_scope env, has_return)
  | Empty -> (env, false)
  | ExprStmt e -> ignore (check_expr env e); (env, false)
  | VarDecl (id, expr) ->
      if Option.is_some (find_var env id) then
        raise (Semantic_error ("Duplicate variable declaration: " ^ id));
      let etype = check_expr env expr in
      if etype != Int then
        raise (Semantic_error "Variable must be initialized with integer");
      (add_var env id { vtype = Int; defined_at = None }, false)
  | Assign (id, expr) ->
      if Option.is_none (find_var env id) then
        raise (Semantic_error ("Assignment to undefined variable: " ^ id));
      let etype = check_expr env expr in
      if etype != Int then
        raise (Semantic_error "Assignment must be to integer variable");
      (env, false)
  | If (cond, then_branch, else_branch) ->
      let ctype = check_expr env cond in
      if ctype != Int then
        raise (Semantic_error "Condition must be integer");
      let env, then_has_return = check_stmt env then_branch return_type in
      let env, else_has_return = 
        match else_branch with
        | Some stmt -> check_stmt env stmt return_type
        | None -> (env, false)
      in
      (env, then_has_return && else_has_return) (* 只有两个分支都有return才算有返回 *)
  | While (cond, body) ->
      let ctype = check_expr env cond in
      if ctype != Int then
        raise (Semantic_error "While condition must be integer");
      let loop_env = { env with in_loop = true } in
      let _, body_has_return = check_stmt loop_env body return_type in
      (env, body_has_return) (* 循环体内的return不影响函数返回状态 *)
  | Break | Continue ->
      if not env.in_loop then
        raise (Semantic_error "Break/continue outside loop");
      (env, false)
  | Return expr_opt ->
      match expr_opt, return_type with
      | Some expr, Int -> 
          let etype = check_expr env expr in
          if etype != Int then
            raise (Semantic_error "Return value must be integer");
          (env, true) (* 标记有return语句 *)
      | None, Void -> (env, true) (* 标记有return语句 *)
      | Some _, Void -> 
          raise (Semantic_error "Void function cannot return value")
      | None, Int -> 
          raise (Semantic_error "Int function must return value")

(* 检查函数是否有return语句 *)
let check_function func env =
  if StringMap.mem func.name env.funcs then
    raise (Semantic_error ("Duplicate function: " ^ func.name));
  
  (* 创建函数作用域 *)
  let env = enter_scope env in
  let env = 
    List.fold_left (fun env (ptype, pname) ->
      if Option.is_some (find_var env pname) then
        raise (Semantic_error ("Duplicate parameter: " ^ pname));
      add_var env pname { vtype = ptype; defined_at = None }
    ) env func.params
  in
  
  let func_env = { 
    env with 
    in_loop = false;
    current_func = Some func.name (* 设置当前函数名 *)
  } in
  
  (* 检查函数体并获取是否有return语句 *)
  let func_env, has_return = 
    List.fold_left (fun (env, has_ret) stmt -> 
        let env, ret = check_stmt env stmt func.return_type in
        (env, has_ret || ret)
      ) (func_env, false) func.body
  in
  
  (* 离开函数作用域 *)
  let env = leave_scope func_env in
  
  (* 非void函数必须有return语句 *)
  if func.return_type = Int && not has_return then
    raise (Semantic_error (Printf.sprintf 
             "Function '%s' must have a return statement" func.name));
  
  { env with funcs = StringMap.add func.name {
    return_type = func.return_type;
    params = func.params;
    defined_at = None;
    has_return = has_return; (* 记录是否有return *)
  } env.funcs }

(* 修改：检查main函数是否有return语句 *)
let check_program program =
  let env = List.fold_left (fun env func -> 
    check_function func env
  ) empty_env program in
  
  if not (StringMap.mem "main" env.funcs) then
    raise (Semantic_error "Missing main function");
  
  let main_func = StringMap.find "main" env.funcs in
  if main_func.return_type != Int then
    raise (Semantic_error "Main must return int");
  if main_func.params != [] then
    raise (Semantic_error "Main must have no parameters");
  
  (* 检查main函数是否有return语句 *)
  if not main_func.has_return then
    raise (Semantic_error "Main function must have a return statement");
  
  env